import numpy as np
import matplotlib.pyplot as plt

# 初始化栅格地图
grid_size = 30
grid = np.full((grid_size, grid_size), 0.8)  # 初始占用概率为 0.5

# 定义传感器参数
# P_Z_given_O = 0.6  # 在网格被占用的情况下检测到障碍物的概率
# P_Z_given_not_O = 0.4  # 在网格未被占用的情况下错误检测到障碍物的概率

P_Z_given_O = 0.2  # 在网格被占用的情况下检测到障碍物的概率
P_Z_given_not_O = 0.8  # 在网格未被占用的情况下错误检测到障碍物的概率

# 定义更新函数
def update_probability(P_O, P_Z_given_O, P_Z_given_not_O):
    P_Z = P_Z_given_O * P_O + P_Z_given_not_O * (1 - P_O)
    P_O_given_Z = (P_Z_given_O * P_O) / P_Z
    print(P_O_given_Z)
    return P_O_given_Z

# 进行观测并更新中心网格
num_observations = 5
center = (grid_size // 2, grid_size // 2)

for i in range(num_observations):
    grid[center] = update_probability(grid[center], P_Z_given_O, P_Z_given_not_O)
    
    
    # 可视化并保存结果
    plt.imshow(grid, cmap='viridis', vmin=0, vmax=1)
    plt.colorbar(label='Occupancy Probability')
    plt.title(f'Grid Map After Observation {i + 1}')
    plt.xlabel('X')
    plt.ylabel('Y')
    
    # 保存图像到文件
    # plt.savefig(f'grid_map_observation_{i + 1}.png')  # 保存为 'grid_map_observation_1.png', 'grid_map_observation_2.png', ...
    plt.close()  # 关闭当前图像，避免重叠
